# -*- coding: utf-8 -*-
from flask import Flask, render_template, request, jsonify
import os
import requests
import re
import json
import oss2
from auth import auth_login, require_login

app = Flask(__name__)

endpoint = f"https://{os.environ['ENDPOINT']}"
auth = oss2.Auth(os.environ['ALIBABA_CLOUD_ACCESS_KEY_ID'], os.environ['ALIBABA_CLOUD_ACCESS_KEY_SECRET'])
bucket = oss2.Bucket(auth, endpoint, os.environ['OSS_BUCKET'])

app.config['UPLOAD_FOLDER'] = './'
app.config['ALLOWED_EXTENSIONS'] = {'jpeg', 'png', 'jpg'}


def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


@app.route('/login', methods=['POST', 'GET'])
def login():
    """登录接口"""
    return auth_login()


@app.route('/', methods=['GET'])
@require_login
def index():
    return render_template('index.html')


@app.route('/api/upload', methods=['POST'])
@require_login
def handle_upload():
    """上传图片接口"""
    try:
        if 'file' not in request.files:
            return 'No file part'
        file = request.files['file']
        if file.filename == '':
            return jsonify({"error": "No selected file"}), 400
        if file and allowed_file(file.filename):
            result = bucket.put_object(file.filename, file.stream)
            if result.status == 200:
                oss_sign_url = bucket.sign_url('GET', file.filename, 3600)
                data = {'url': oss_sign_url}
                return jsonify({'code': 200, 'data': data})
            else:
                return jsonify({'code': 200, 'data': {}})
        else:
            return jsonify({"error": "File type not allowed"}), 409
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/completions', methods=['POST'])
@require_login
def handle_completions():
    """信息提取接口"""
    try:
        data = request.json
        oss_url = data.get('url')
        prompt = data.get('prompt', '')
        response = curl_request(oss_url, prompt)
        data = None
        if (response.status_code == 200):
            resp_content = response.json()['choices'][0]['message']['content'];
            data = extract_code_blocks(resp_content)
            data = json.loads(data)
        if data is None:
            return jsonify({'code': 400, 'data': {}})
        return jsonify({'code': 200, 'data': data, 'src': oss_url})
    except Exception as e:
        return jsonify({'error': str(e)}), 500


def curl_request(oss_url, extra_info):
    """调用百炼模型服务"""
    prompt = (
        "# 角色\n"
        "你是一名图片分析师，擅长对各种图片进行理解。\n\n"
        "# 任务描述\n"
        "给你一个张图片，请你完成以下任务。\n"
        "- 分析图片内容，提取关键信息。\n"
        "# 限制\n"
        "- 分析范围严格限定于提供的图片内容，不涉及图片之外的任何推测或其他信息。\n"
        "- 总结时需严格依据图片内容，不可添加个人臆测或创意性内容。\n"
        "- 只能使用JSON返回结果，key必须使用中文，JSON数据格式必须正确。\n"
        "# 输入数据\n"
        "## 图片提取内容补充信息 (你需要针对补充信息，进行图片内容提取，如果输入为空则忽略补充信息)\n"
        f"{extra_info}\n\n"
        "# 输出格式\n"
        "直接JSON格式输出内容。\n"
    )
    api_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions'
    # 设置请求头
    headers = {
        'Authorization': 'Bearer ' + os.environ['DASHSCOPE_API_KEY'],
        'Content-Type': 'application/json; charset=utf-8'
    }
    data = {
        'model': 'qwen-vl-max',
        'messages': [
            {
                'role': 'user',
                'content': [
                    {
                        'type': 'image_url',
                        'image_url': {
                            'url': oss_url
                        }
                    },
                    {
                        'type': 'text',
                        'text': prompt
                    }
                ]
            }
        ]
    }
    post_data = json.dumps(data, ensure_ascii=False, indent=4)
    # 发送请求
    return requests.post(api_url, data=post_data.encode('utf-8'), headers=headers)


def extract_code_blocks(text):
    """ 正则表达式匹配  json包裹的代码块 """
    pattern = r'```json\n(.*?)```'
    # 使用re.findall()函数查找所有匹配的代码块
    code_blocks = re.findall(pattern, text, re.DOTALL)

    if len(code_blocks) > 0:
        return code_blocks[0]

    return ''


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=9000)
